"""
# 不同方法生成训练集和测试集
# 按比例划分训练集和测试集
# 五折交叉验证训练集和测试集

"""
import os
import csv
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split, StratifiedKFold


# 为二维模型生成，第一种，直接按比例划分训练集和测试集，不考虑数据类别不平衡问题
def generate_one():
    path = '/data1/zmy/data_augmentation/all_data3.csv'
    data = []
    f = csv.reader(open(path, 'r'))
    for i in f:
        data.append(i)
    
    labels = []
    features = []
    for i in range(1, len(data)):
        features.append(data[i])
        labels.append(int(data[i][5]))

    features = np.array(features)
    labels = np.array(labels, dtype=np.int)
    x_train, x_test, y_train, y_test = train_test_split(features,
                                                        labels,
                                                        test_size=0.2, 
                                                        random_state=1234565, 
                                                        stratify=labels)

    save_path = '/data1/zmy/data_augmentation/divide_csv/generate_one/'
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    df = pd.DataFrame(x_train, columns=data[0])
    df.to_csv(save_path+'train.csv', index=False)

    df2 = pd.DataFrame(x_test, columns=data[0])
    df2.to_csv(save_path+'test.csv', index=False)


# 为三维模型生成，第一种，生成五折交叉训练集和测试集，不考虑数据类别不平衡问题
def generate_two():
    path = '/data1/zmy/data_augmentation/labels.csv'
    data = []
    f = csv.reader(open(path, 'r'))
    for i in f:
        data.append(i)

    labels = []
    features = []
    for i in range(1, len(data)):
        features.append(data[i])
        labels.append(int(data[i][6]))

    features = np.array(features)
    labels = np.array(labels, dtype=np.int)

    # 把所有的测试集保存下来，就是相互独立的五折数据集
    sfolder = StratifiedKFold(n_splits=5, random_state=0)

    # 定义五折数据集顺序和csv文件列名
    index = 0
    row_name = data[0]

    # 生成存储文件夹
    save_path = '/data1/zmy/data_augmentation/kfold_dataset/'
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    for train, test in sfolder.split(features, labels):
        print('train: %d | test: %d' % (len(train), len(test)))
        son_features = features[test]
        df = pd.DataFrame(son_features, columns=row_name)
        df.to_csv(save_path + str(index) + '.csv', index=False)
        index = index + 1


# 为三维模型生成， 第二种，直接按比例划分训练集和测试集，不考虑数据类别不平衡问题
def generate_three():
    path = '/data1/zmy/data_augmentation/labels.csv'
    data = []
    f = csv.reader(open(path, 'r'))
    for i in f:
        data.append(i)

    labels = []
    features = []
    for i in range(1, len(data)):
        features.append(data[i])
        labels.append(int(data[i][6]))

    features = np.array(features)
    labels = np.array(labels, dtype=np.int)
    x_train, x_test, y_train, y_test = train_test_split(features,
                                                        labels,
                                                        test_size=0.2,
                                                        random_state=1234565,
                                                        stratify=labels)

    save_path = '/data1/zmy/data_augmentation/divide_csv/generate_three/'
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    df = pd.DataFrame(x_train, columns=data[0])
    df.to_csv(save_path + 'train.csv', index=False)

    df2 = pd.DataFrame(x_test, columns=data[0])
    df2.to_csv(save_path + 'test.csv', index=False)


if __name__ == '__main__':
    # generate_one()
    # generate_two()
    generate_three()